import subprocess
import time
import threading
import logging
from prometheus_client import start_http_server, Gauge
import json

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

# List of RTMP streams to monitor
STREAMS = [
    "rtmp://nginx-rtmp:1935/live/stream"
]

# Probe interval
PROBE_INTERVAL = 60

# Prometheus metrics
stream_up = Gauge("ffmpeg_stream_up", "Stream availability", ["stream"])
frame_rate = Gauge("ffmpeg_stream_frame_rate", "Frame rate of video stream", ["stream"])
frame_width = Gauge("ffmpeg_stream_frame_width", "Width of video stream", ["stream"])
frame_height = Gauge("ffmpeg_stream_frame_height", "Height of video stream", ["stream"])
video_codec_present = Gauge("ffmpeg_stream_video_codec", "Presence of video codec", ["stream", "codec"])
audio_codec_present = Gauge("ffmpeg_stream_audio_codec", "Presence of audio codec", ["stream", "codec"])


def probe_stream(url):
    logging.info(f"Probing stream: {url}")
    try:
        result = subprocess.run([
            "ffprobe",
            "-v", "error",
            "-show_streams",
            "-of", "json",
            url
        ], capture_output=True, text=True, timeout=60)

        if result.returncode != 0:
            logging.warning(f"Stream {url} not available (ffprobe error).")
            stream_up.labels(stream=url).set(0)
            frame_rate.labels(stream=url).set(0)
            frame_width.labels(stream=url).set(0)
            frame_height.labels(stream=url).set(0)
            return

        data = json.loads(result.stdout)
        video_found = False
        audio_found = False

        for stream in data.get("streams", []):
            if stream.get("codec_type") == "video":
                video_found = True
                codec = stream.get("codec_name", "unknown")
                width = stream.get("width", 0)
                height = stream.get("height", 0)
                avg_frame_rate = stream.get("avg_frame_rate", "0/1")
                try:
                    num, denom = map(int, avg_frame_rate.split("/"))
                    rate = num / denom if denom else 0.0
                except:
                    rate = 0.0

                frame_rate.labels(stream=url).set(rate)
                frame_width.labels(stream=url).set(width)
                frame_height.labels(stream=url).set(height)
                video_codec_present.labels(stream=url, codec=codec).set(1)

            elif stream.get("codec_type") == "audio":
                audio_found = True
                codec = stream.get("codec_name", "unknown")
                audio_codec_present.labels(stream=url, codec=codec).set(1)

        # Reset metrics for codecs not seen
        if not video_found:
            frame_rate.labels(stream=url).set(0)
            frame_width.labels(stream=url).set(0)
            frame_height.labels(stream=url).set(0)
        if not audio_found:
            audio_codec_present.labels(stream=url, codec="unknown").set(0)

        stream_up.labels(stream=url).set(1 if video_found or audio_found else 0)

    except subprocess.TimeoutExpired:
        logging.error(f"ffprobe timed out on stream {url}")
        stream_up.labels(stream=url).set(0)
    except Exception as e:
        logging.error(f"Unexpected error probing {url}: {e}")
        stream_up.labels(stream=url).set(0)


def run_probe_loop():
    while True:
        for url in STREAMS:
            probe_stream(url)
        time.sleep(PROBE_INTERVAL)


if __name__ == "__main__":
    start_http_server(8001)
    logging.info("Starting ffmpeg probe exporter on port 8001")
    threading.Thread(target=run_probe_loop, daemon=True).start()
    while True:
        time.sleep(5)  # Keep main thread alive